package main

import (
	"fmt"
	"math"
)

// https://leetcode.cn/problems/best-time-to-buy-and-sell-stock-ii/?envType=study-plan-v2&envId=top-interview-150
// 给你一个整数数组 prices ，其中 prices[i] 表示某支股票第 i 天的价格。
//在每一天，你可以决定是否购买和/或出售股票。你在任何时候 最多 只能持有 一股 股票。你也可以先购买，然后在 同一天 出售。
//返回 你能获得的 最大 利润 。

func main() {
	prices := []int{7, 1, 5, 3, 6, 4}
	val := maxProfitSimple(prices)
	fmt.Println(val)
}

// 简单版本
// dp[i][0] = max(dp[i-1][0], dp[i-1][1] + prices[i])
// dp[i][1] = max(dp[i-1][1], dp[i-1][0] - prices[i])
func maxProfitSimple(prices []int) int {
	dayNum := len(prices)
	dp := map[int]map[int]int{}
	for i := 0; i < dayNum; i++ {
		if dp[i] == nil {
			dp[i] = map[int]int{}
		}
		if i == 0 {
			dp[i][0] = 0
			dp[i][1] = -prices[i]
			continue
		}
		dp[i][0] = max(dp[i-1][0], dp[i-1][1]+prices[i])
		dp[i][1] = max(dp[i-1][1], dp[i-1][0]-prices[i])
	}
	return max(dp[dayNum-1][0], dp[dayNum-1][1])
}

func max(v1, v2 int) int {
	if v1 > v2 {
		return v1
	}
	return v2
}

// 通用版本
// dp[i][k][0] = max(dp[i-1][k][0], dp[i-1][k][1] + prices[i])
// dp[i][k][1] = max(dp[i-1][k][1], dp[i-1][k-1][0] - prices[i])
// dp[0][...][1] = prices[0]
// dp[0][...][0] = 0
func maxProfit(prices []int) int {
	dp := make(map[int]map[int]map[int]int) // 第x天 k次交易 有无股票 的利润
	maxK := 2
	dayNum := len(prices)
	dp[0] = map[int]map[int]int{}

	for i := 0; i < dayNum; i++ {
		for j := 0; j < maxK; j++ {
			if i == 0 {
				dp = setVal(dp, 0, i, j, false) //dp[i][j][0] = 0
				if j == 0 {
					dp = setVal(dp, math.MinInt, i, j, true) // dp[i][j][1] = math.MinInt
				} else {
					dp = setVal(dp, -prices[i], i, j, true) // dp[i][j][1] = -prices[i]
				}
			}

			// dp[i][j][0] = max(dp[i-1][j][0], dp[i-1][j][1]+prices[i])
			dp = setVal(dp, max(dp[i-1][j][0], dp[i-1][j][1]+prices[i]), i, j, false)
			// dp[i][j][1] = max(dp[i-1][j][1], dp[i-1][j-1][0]-prices[i])
			dp = setVal(dp, max(dp[i-1][j][1], dp[i-1][j-1][0]-prices[i]), i, j, true)
		}
	}
	return max(getVal(dp, dayNum-1, 0, false), getVal(dp, dayNum-1, 0, true))
}

func setVal(dp map[int]map[int]map[int]int, val, i, k int, has bool) map[int]map[int]map[int]int {
	if dp[i] == nil {
		dp[i] = map[int]map[int]int{}
	}
	if dp[i][k] == nil {
		dp[i][k] = map[int]int{}
	}
	vHas := 0
	if has {
		vHas = 1
	}
	dp[i][k][vHas] = val
	return dp
}

func getVal(dp map[int]map[int]map[int]int, i, k int, has bool) int {
	if dp[i] == nil || dp[i][k] == nil {
		return 0
	}
	vHas := 0
	if has {
		vHas = 1
	}
	return dp[i][k][vHas]
}
